package org.znxs.znmanus.advisors;

import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.model.MessageAggregator;
import reactor.core.publisher.Flux;

import java.util.List;

/**
 * 自定义打印日志 Advisor
 * 打印 info 级别的日志 、 只输出单词用户提示词和AI回复的文本
 */
@Slf4j
public class MyLoggerAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {

    /**
     * 输出ai响应的内容 后置环绕
     *
     * @param advisedResponse advisor 返回结果
     */
    private void observeAfter(AdvisedResponse advisedResponse) {
        log.info("\n AI Response: {}", advisedResponse.response().getResult().getOutput().getText());
        // 如果调用了工具，显示所有调用了的工具
        if (advisedResponse.response().hasToolCalls()) {
            // 调用了工具
            List<AssistantMessage.ToolCall> toolCalls =
                    advisedResponse.response().getResult().getOutput().getToolCalls();
            toolCalls.forEach(tc -> log.info("\nAI 使用了工具 :{} 调用的参数为:{}", tc.name(), tc.arguments()));
        }


    }

    /**
     * 输入用户的提问 前置环绕
     *
     * @param advisedRequest advisor 请求封装
     * @return
     */
    private AdvisedRequest before(AdvisedRequest advisedRequest) {
        log.info("\n AI Request: {}", advisedRequest.userText());
        // 如果有工具 这里显示所有可调用工具
        return advisedRequest;
    }

    @Override
    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
        advisedRequest = this.before(advisedRequest);
        AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);
        this.observeAfter(advisedResponse);
        return advisedResponse;
    }

    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
        advisedRequest = this.before(advisedRequest);
        Flux<AdvisedResponse> advisedResponseFlux = chain.nextAroundStream(advisedRequest);
        //通过 MessageAggregator 工具类将 Flux 响应聚合成单个 AdvisedResponse。
        // 这对于日志记录或其他需要观察整个响应而非流中各个独立项的处理非常有用。
        // 不过不能在这里修改响应，因为是该响应是只读的
        return (new MessageAggregator()).aggregateAdvisedResponse(advisedResponseFlux, this::observeAfter);
    }

    @NotNull
    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }

    @Override
    public int getOrder() {
        return 0;
    }
}
